# test/postgresql/conftest.py
"""
postgresql 的测试环境配置
"""

import os
from dotenv import load_dotenv

import pytest
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from models.user import Base

load_dotenv()

postgres_host = os.getenv("POSTGRES_HOST")
postgres_port = os.getenv("POSTGRES_PORT")
postgres_db = os.getenv("POSTGRES_DB")
postgres_user = os.getenv("POSTGRES_USER")
postgres_password = os.getenv("POSTGRES_PASSWORD")

# 创建一个不带数据库名的URL，用于连接默认数据库
DEFAULT_DB_URL = f"postgresql://{postgres_user}:{postgres_password}@{postgres_host}:{postgres_port}/postgres"
DATABASE_URL = f"postgresql://{postgres_user}:{postgres_password}@{postgres_host}:{postgres_port}/{postgres_db}"


def ensure_database_exists():
    """确保测试数据库存在，如果不存在则创建"""
    # 连接到默认的postgres数据库
    default_engine = create_engine(DEFAULT_DB_URL)

    try:
        # 检查数据库是否存在
        with default_engine.connect() as conn:
            result = conn.execute(text(f"SELECT 1 FROM pg_database WHERE datname = '{postgres_db}'"))
            database_exists = result.scalar() is not None

            if not database_exists:
                # 如果数据库不存在，创建它
                # 需要在事务外执行CREATE DATABASE
                conn.execute(text("COMMIT"))
                conn.execute(text(f"CREATE DATABASE {postgres_db}"))
                print(f"Created database {postgres_db}")
    finally:
        default_engine.dispose()


# 创建数据库引擎
engine = create_engine(DATABASE_URL, echo=True)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def create_tables():
    """创建所有表（如果不存在）"""
    Base.metadata.create_all(bind=engine)


# pytest fixture：为每个测试提供一个新的数据库会话
@pytest.fixture(scope="function")
def db_session():
    ensure_database_exists()  # 确保数据库存在
    create_tables()  # 创建表（如果不存在）
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()
